# coding = utf-8

'''
针对mU-Net的一些相关可视化的工作
'''
import os
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import cv2

IMAGE_PATH = "/datasets/DongbeiDaxue/chengkun"
OUTPUT_PATH = "/home/diaozhaoshuo/log/BeliefFunctionNN/chengkung/dongbeidaxue/munet/image_tumor"

def show_data(case_id):
   origion_data_path = os.path.join(IMAGE_PATH, "case_{}".format(str(case_id).zfill(5)))
   origion_data_path = os.path.join(origion_data_path, "imaging")

   label_data_path = os.path.join(OUTPUT_PATH, "case_{}".format(str(case_id).zfill(5)))
   label_data_path = os.path.join(label_data_path, "label")

   predict_data_path = os.path.join(OUTPUT_PATH, "case_{}".format(str(case_id).zfill(5)))
   predict_data_path = os.path.join(predict_data_path, "predict")

   for i in range(len(os.listdir(origion_data_path))):
       origion_file = os.path.join(origion_data_path, "{}.npy".format(str(i).zfill(3)))
       label_file = os.path.join(label_data_path, "{}.png".format(str(i).zfill(3)))
       predict_file = os.path.join(predict_data_path, "{}.png".format(str(i).zfill(3)))

       image = np.load(origion_file)
       image = image * 255
       image = image.astype(np.uint8)

       image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)

       label = Image.open(label_file).convert("L")
       label = np.asarray(label)
       output = Image.open(predict_file).convert("L")
       output = np.asarray(output)
       if np.max(label) == 0 and np.max(output) == 0:
           continue



       contours, _ = cv2.findContours(label, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
       for counter in contours:
           data_list = []
           for i in range(counter.shape[0]):
               j = counter[i][0]
               data_list.append(j)
           cv2.polylines(image, np.array([data_list], np.int32), True,[0,255, 0], thickness=1)

       contours, _ = cv2.findContours(output, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

       for counter in contours:
           data_list = []
           for i in range(counter.shape[0]):
               j = counter[i][0]
               data_list.append(j)
           cv2.polylines(image, np.array([data_list], np.int32), True, [0, 0, 255], thickness=1)

       plt.imshow(image)
       plt.show()

def show_image():
    origion_image = os.path.join(IMAGE_PATH, "case_00035/imaging/766.npy")
    fusion_image = os.path.join(OUTPUT_PATH, "case_00035/image/767.png")
    origion_image = np.load(origion_image)
    origion_image = origion_image * 255
    origion_image = origion_image.astype(np.uint8)
    origion_image = cv2.cvtColor(origion_image, cv2.COLOR_GRAY2BGR)
    fusion_image = Image.open(fusion_image)
    plt.subplot(1, 2, 1)
    plt.imshow(origion_image)
    plt.subplot(1, 2, 2)
    plt.imshow(fusion_image)
    plt.show()


if __name__ == '__main__':
    show_image()